import os
import time
from typing import List
from dotenv import load_dotenv

from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import Runnable, RunnablePassthrough
from langchain_core.documents import Document
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_openai import ChatOpenAI
from langchain.text_splitter import RecursiveCharacterTextSplitter
from pdfminer.high_level import extract_pages
from pdfminer.layout import LTTextContainer

# 文本分割和速率限制的常量
CHUNK_SIZE = 3000
CHUNK_OVERLAP = 200

# Markdown 转换提示模板
MARKDOWN_PROMPT_TEMPLATE = """
**说明：**

*   使用适当代表文本内容的标题（#、##、###等）。
*   使用项目符号或编号列表有效地组织信息。
*   使用**粗体**或*斜体*强调重要术语或短语。
*   使用代码块（```）表示代码（如需要）。
*   使用表格（|）组织信息。
*   删除不必要的换行或额外空格。
*   遵循标准的 Markdown 语法。

**文本：**
{text}
"""

def extract_text_from_pdf(pdf_path: str) -> List[str]:
    """
    逐页从 PDF 文件中提取文本。

    参数：
        pdf_path: PDF 文件的路径。

    返回：
        字符串列表，每个字符串代表一页的文本内容。
    """
    all_texts = []
    for page_layout in extract_pages(pdf_path):
        page_text = "".join(
            element.get_text()
            for element in page_layout
            if isinstance(element, LTTextContainer)
        )
        all_texts.append(page_text)
    return all_texts

def create_markdown_chain(llm: Runnable) -> Runnable:
    """
    创建用于将文本转换为 Markdown 的 LangChain 链。

    参数：
      llm: 要使用的语言模型可运行对象（例如 ChatGoogleGenerativeAI、ChatOpenAI）。

    返回：
      LangChain 链（可运行对象）。
    """
    prompt = PromptTemplate.from_template(MARKDOWN_PROMPT_TEMPLATE)
    return (
        {"text": RunnablePassthrough()}
        | prompt
        | llm
        | StrOutputParser()
    )

def convert_to_markdown(text: str, llm: Runnable) -> str:
    """
    使用 LangChain 和指定的大语言模型将文本转换为 Markdown 格式。

    参数：
        text: 要转换的文本。
        llm: 语言模型可运行对象实例。

    返回：
        转换为 Markdown 格式的文本。
    """
    chain = create_markdown_chain(llm)
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP, add_start_index=True
    )

    docs = text_splitter.split_documents([Document(page_content=text)])
    markdown_text = ""
    for doc in docs:
        markdown_text += chain.invoke(doc.page_content) + "\n"
    return markdown_text

def save_markdown_to_file(markdown_texts: List[str], output_filename: str = "output.md"):
    """保存生成的 Markdown 到文件。"""
    with open(output_filename, "w", encoding="utf-8") as f:
        for page in markdown_texts:
            f.write(page)
            f.write("\n")

def main():
    """运行 PDF 到 Markdown 转换的主函数。"""
    load_dotenv()
    pdf_path = os.getenv("PDF_PATH", "test.pdf")
    
    # 配置大语言模型提供商和 API 密钥
    llm_provider = os.getenv("LLM_PROVIDER", "openai")
    google_api_key = os.getenv("GOOGLE_API_KEY")
    openai_api_key = os.getenv("OPENAI_API_KEY")
    openai_base_url = os.getenv("OPENAI_BASE_URL")
    openai_model_name = os.getenv("OPENAI_MODEL_NAME", "gpt-4o-mini")
    gemini_model_name = os.getenv("GEMINI_MODEL_NAME", "gemini-1.5-flash")

    # Gemini
    if llm_provider == "google" and not google_api_key:
        raise ValueError(
            "环境变量 'GOOGLE_API_KEY' 未设置。"
            "请设置以使用 Google Gemini 模型。"
        )
    # GPT
    elif llm_provider == "openai" and not openai_api_key:
        raise ValueError(
            "环境变量 'OPENAI_API_KEY' 未设置。"
            "请设置以使用 OpenAI 模型。"
        )

    # 根据选定的提供商创建大语言模型实例
    if llm_provider == "google":
        llm = ChatGoogleGenerativeAI(model=gemini_model_name, google_api_key=google_api_key, convert_system_message_to_human=True, temperature=0.5) 
    elif llm_provider == "openai":
        llm = ChatOpenAI(api_key=openai_api_key, model=openai_model_name, temperature=0.7, base_url=openai_base_url)
    else:
        raise ValueError(f"无效的大语言模型提供商：{llm_provider}")

    extracted_texts = extract_text_from_pdf(pdf_path)

    markdown_output = []
    for text in extracted_texts:
        markdown_text = convert_to_markdown(text, llm)
        markdown_output.append(markdown_text)

    for i, page in enumerate(markdown_output):
        print(f"--- 第 {i+1} 页 ---")
        print(page)

    save_markdown_to_file(markdown_output)
    print(f"Markdown 文件已保存到 'output.md'。")

if __name__ == "__main__":
    main()